# middleware/unified_response.py
from fastapi import Request, Response, FastAPI
from starlette.middleware.base import BaseHTTPMiddleware

class UnifiedResponseMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        response = await call_next(request)
        if response.status_code < 400:
            # 假设响应体是 JSON 格式的
            response_body = [chunk async for chunk in response.body_iterator]
            response_body_str = response_body[0].decode()
            response_body_dict = response.json()
            
            # 统一格式化响应
            unified_response = {
                "code": 0,
                "message": "success",
                "data": response_body_dict
            }
            response.body = json.dumps(unified_response).encode('utf-8')
            response.media_type = "application/json"
        return response